from PySide6.QtGui import QPixmap
from PySide6.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QFrame, QPushButton, QFileDialog
from utils.detect_api import Detector


class MyWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("垃圾分类器")
        self.resize(1200, 600)
        # 原始图像路径
        self.image_path = ""
        self.detector = Detector()

        # 创建三个垂直布局
        layout = QHBoxLayout(self)
        layout.setContentsMargins(10, 10, 10, 10)

        left_layout = QVBoxLayout()
        center_layout = QVBoxLayout()
        right_layout = QVBoxLayout()

        # 创建左侧布局的预览区域
        self.image_label = QLabel(self)
        self.image_label.setFixedSize(400, 400)  # 设置固定大小的预览区域

        # 在中间布局中添加选择图像按钮和开始识别按钮
        select_image_button = QPushButton("选择图像")
        select_image_button.clicked.connect(self.select_image)

        detect_img_by_vgg16_btn = QPushButton("vgg16识别")
        detect_img_by_vgg16_btn.clicked.connect(self.detect_img_by_vgg16)

        detect_img_by_resnet50_btn = QPushButton("ResNet50识别")
        detect_img_by_resnet50_btn.clicked.connect(self.detect_img_by_resnet50_btn)

        # 右侧布局显示结果文本
        self.result_img_label = QLabel(self)
        self.result_img_label.setFixedSize(400, 400)

        # 添加到右侧布局中
        right_layout.addWidget(self.result_img_label)

        # 创建分割线
        separator1 = QFrame()
        separator1.setFrameShape(QFrame.VLine)
        separator1.setFrameShadow(QFrame.Sunken)

        separator2 = QFrame()
        separator2.setFrameShape(QFrame.VLine)
        separator2.setFrameShadow(QFrame.Sunken)

        # 将预览区域和其他布局组合到主布局中
        left_layout.addWidget(self.image_label)
        layout.addLayout(left_layout)
        layout.addWidget(separator1)
        layout.addLayout(center_layout)
        layout.addWidget(separator2)
        layout.addLayout(right_layout)

        # 在中间布局中添加选择图像按钮和开始识别按钮
        center_layout.addStretch()  # 添加弹簧以创建垂直居中的效果
        center_layout.addWidget(select_image_button)
        center_layout.addWidget(detect_img_by_vgg16_btn)
        center_layout.addWidget(detect_img_by_resnet50_btn)
        center_layout.addStretch()  # 添加弹簧以创建垂直居中的效果

    # 绑定选择图片槽函数
    def select_image(self):
        options = QFileDialog.Options()
        options |= QFileDialog.DontUseNativeDialog  # 禁用原生文件对话框（可选）

        file_path, _ = QFileDialog.getOpenFileName(self, "选择图像", "", "图像文件 (*.png *.jpg *.jpeg)",
                                                   options=options)
        self.image_path = file_path

        if file_path:
            pixmap = QPixmap(file_path)
            self.image_label.setPixmap(pixmap)

    # vgg16识别按钮
    def detect_img_by_vgg16(self):
        pred_str = self.detector.vgg16_predict(image_path=self.image_path)
        self.result_img_label.setText(pred_str)

    # resnet50识别按钮
    def detect_img_by_resnet50_btn(self):
        pred_str = self.detector.resnet50_predict(image_path=self.image_path)
        self.result_img_label.setText(pred_str)


if __name__ == "__main__":
    app = QApplication([])
    window = MyWindow()
    window.show()
    app.exec()
